#!/usr/bin/python3
import sys
sys.path.insert(0, '.')
import numpy as np
import argparse
import os
import cv2
from aip import AipFace
import base64
import time
import json

""" 你的 APPID AK SK """
APP_ID = '16315779'
API_KEY = 'NqzAMufM5INE9AsSS8yw9bek'
SECRET_KEY = 'HtLXQXxdTz9gLL1hGrQ7p1y7oDcTMnN9'

SHOW = False
# SHOW = True

MAX_OK_IMAGE_COUNT = 1000

IMAGE_PATH = r'/rootfs/media/yery/Data1/data/VideoCropFace'
JSON_PATH = r'/rootfs/media/yery/Data1/data/VideoCropFace/FaceAttrJson'
IMAGE_LIST_FILES = [
    # r'/rootfs/media/yery/Data1/data/VideoCropFace/face_image_q5-10.txt',
    r'/rootfs/media/yery/Data1/data/VideoCropFace/face_image_q10-15.txt',
    r'/rootfs/media/yery/Data1/data/VideoCropFace/face_image_q15-20.txt',
    r'/rootfs/media/yery/Data1/data/VideoCropFace/face_image_q20-25.txt',
    r'/rootfs/media/yery/Data1/data/VideoCropFace/face_image_q25-30.txt',
    r'/rootfs/media/yery/Data1/data/VideoCropFace/face_image_q30-35.txt',
    r'/rootfs/media/yery/Data1/data/VideoCropFace/face_image_q35-40.txt',
    r'/rootfs/media/yery/Data1/data/VideoCropFace/face_image_q40-45.txt',
    r'/rootfs/media/yery/Data1/data/VideoCropFace/face_image_q45-50.txt',
    r'/rootfs/media/yery/Data1/data/VideoCropFace/face_image_q50-55.txt',
    r'/rootfs/media/yery/Data1/data/VideoCropFace/face_image_q55-60.txt',
    r'/rootfs/media/yery/Data1/data/VideoCropFace/face_image_q60-65.txt',
    r'/rootfs/media/yery/Data1/data/VideoCropFace/face_image_q65-70.txt',
    r'/rootfs/media/yery/Data1/data/VideoCropFace/face_image_q70-75.txt',
    r'/rootfs/media/yery/Data1/data/VideoCropFace/face_image_q75-80.txt',
    r'/rootfs/media/yery/Data1/data/VideoCropFace/face_image_q80-85.txt',
    r'/rootfs/media/yery/Data1/data/VideoCropFace/face_image_q85-90.txt',
    r'/rootfs/media/yery/Data1/data/VideoCropFace/face_image_q90-95.txt',
    r'/rootfs/media/yery/Data1/data/VideoCropFace/face_image_q95-100.txt',
]


""" 读取图片 """
def get_file_content(filePath):
    with open(filePath, 'rb') as fp:
        return str(base64.b64encode(fp.read()), 'utf-8')


def conv_to_base64(image_jpg_stream):
    return str(base64.b64encode(image_jpg_stream), 'utf-8')


def detect(client, imagepathbase, imagename1, dbg_str='', err_file=None):

    imagepath1 = os.path.join(imagepathbase, imagename1)

    image_file1 = open(imagepath1, 'rb')
    image_jpg_stream1 = image_file1.read()
    imagestream1 = conv_to_base64(image_jpg_stream1)
    image_file1.close()

    options = {}
    options["max_face_num"] = 1
    # options["face_field"] = "gender,age,glasses,quality,eye_status"
    options["face_field"] = "age,beauty,expression,face_shape,gender,glasses,landmark,landmark150,race,quality,eye_status,emotion,face_type,mask"

    error_count = 0
    while True:
        try:
            ret1 = client.detect(imagestream1, 'BASE64', options)
            break
        except Exception as e:
             # print('1', e)
             error_count += 1
             if error_count >= 2:
                 det_err_info = '1 DetErr: {}'.format(imagename1)
                 print(dbg_str + det_err_info)
                 if err_file is not None:
                     err_file.write(imagename1 + '\n')
                 return None
             time.sleep(1.0)
             continue
    try:
        face_num1 = ret1['result']['face_num']
        if face_num1 > 0:
            face_list1 = ret1['result']['face_list']
            face1 = face_list1[0]
            # face_token1 = face1['face_token']
            location1 = face1['location']
        else:
            det_err_info = '3 DetErr: {}'.format(imagename1)
            print(dbg_str + det_err_info)
            if err_file is not None:
                err_file.write(imagename1 + '\n')
            return None
    except Exception as e:
        # print('2', e)
        det_err_info = '2 DetErr: {}'.format(imagename1)
        print(dbg_str + det_err_info)
        if err_file is not None:
            err_file.write(imagename1 + '\n')
        if SHOW:
            image_org1 = cv2.imread(imagepath1)
            cv2.imshow("image1", image_org1)
            cv2.waitKey()
        return None

    # print(dbg_str + imagename1 + ':', face1)

    if SHOW:
        left1, top1, width1, height1 = int(location1['left']), int(location1['top']), int(location1['width']), int(location1['height'])
        image_org1 = cv2.imread(imagepath1)
        color = (0, 255, 0)
        cv2.rectangle(image_org1, (left1, top1), (left1 + width1, top1 + height1), color, thickness=1)
        cv2.imshow("image1", image_org1)
        cv2.waitKey()

    return face1


def main():
    if not os.path.exists(JSON_PATH):
        os.makedirs(JSON_PATH)
    client = AipFace(APP_ID, API_KEY, SECRET_KEY)
    client.setConnectionTimeoutInMillis(10000)
    client.setSocketTimeoutInMillis(10000)

    for IMAGE_LIST_FILE in IMAGE_LIST_FILES:
        base_name = os.path.splitext(os.path.basename(IMAGE_LIST_FILE))[0]
        err_file_path = os.path.join(IMAGE_PATH, '{}_attr_err.txt'.format(base_name))
        ok_file_path = os.path.join(IMAGE_PATH, '{}_attr_ok.txt'.format(base_name))
        ok_file = open(ok_file_path, 'w')
        with open(err_file_path, 'w') as err_file:
            ok_file_count = 0
            err_file_count = 0
            with open(IMAGE_LIST_FILE, 'r') as file:
                image_list = []
                for image in file.readlines():
                    image = image.split()[0]
                    # image = image.split()
                    # image = os.path.join(image[0], image[1])
                    image_list.append(image.strip())
                # image_list.sort()
                image_count = len(image_list)
                start_time = time.time()
                for i, image in enumerate(image_list):
                    time.sleep(0.5)
                    dbg_str = '{}/{} '.format(i, image_count)
                    face_info = detect(client, IMAGE_PATH, image, dbg_str, err_file)
                    if face_info is None:
                        err_file_count += 1
                        continue
                    try:
                        json_path = os.path.join(JSON_PATH, image).replace('.jpg', '.json')
                        json_dir = os.path.dirname(json_path)
                        if not os.path.exists(json_dir):
                            os.makedirs(json_dir)
                        with open(json_path, 'w') as json_file:
                            json.dump(face_info, json_file)
                    except Exception as e:
                        # print('3', e)
                        err_file_count += 1
                        continue
                    ok_file.write(image + '\n')
                    ok_file_count += 1
                    if ok_file_count >= MAX_OK_IMAGE_COUNT:
                        break
                    if (i+1) % 50 == 0:
                        print('Proccess File: {}/{}, {:.2f}% {:.2f}%, Ok {}, Err {} {:.2f}%, Time: {:.3f}s'.format(i+1, image_count, 100*(i+1)/image_count, 100*ok_file_count/MAX_OK_IMAGE_COUNT, ok_file_count, err_file_count, 100*err_file_count/(err_file_count+ok_file_count), time.time()-start_time))
        ok_file.close()
        os.system('chmod a+wr {}'.format(err_file_path))
        os.system('chmod a+wr {}'.format(ok_file_path))


if __name__ == '__main__':
    main()
